#!/bin/bash

# DevDeploy CV API服务动态启动脚本（CCR版本）
# 用法: ./start_cv_service_ccr.sh [--classifier <path>] [--segment <path>] [--port <port>] [--gpu <gpu_ids>]
# 使用基于腾讯云CCR基础镜像构建的镜像

# 解析参数
CLASSIFIER_MODEL_PATH=""
SEGMENT_MODEL_PATH=""
PORT=8796
GPU_IDS_ARG=""  # 命令行指定的GPU ID，例如 "0" 或 "0,1"

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        --classifier)
            CLASSIFIER_MODEL_PATH="$2"
            shift 2
            ;;
        --segment)
            SEGMENT_MODEL_PATH="$2"
            shift 2
            ;;
        --port)
            PORT="$2"
            shift 2
            ;;
        --gpu)
            GPU_IDS_ARG="$2"
            shift 2
            ;;
        *)
            echo "❌ 未知参数: $1"
            echo ""
            echo "用法: $0 [--classifier <path>] [--segment <path>] [--port <port>] [--gpu <gpu_ids>]"
            echo ""
            echo "参数说明:"
            echo "  --classifier <path>  - 分类模型路径（可选）"
            echo "  --segment <path>     - 分割模型路径（可选）"
            echo "  --port <port>        - 服务端口号（可选，默认8796）"
            echo "  --gpu <gpu_ids>      - GPU设备ID，支持单个或多个（可选，默认0）"
            echo "                        例如: --gpu 0 或 --gpu 0,1 或 --gpu \"0,1\""
            echo ""
            echo "注意: 至少需要配置一个模型（分类或分割）"
            echo ""
            echo "示例:"
            echo "  # 只配置分类模型"
            echo "  $0 --classifier ../weights/resnet50_qiu/fullmodel_best.onnx"
            echo ""
            echo "  # 只配置分割模型"
            echo "  $0 --segment ../weights/segment_model.onnx"
            echo ""
            echo "  # 同时配置分类和分割模型"
            echo "  $0 --classifier ../weights/resnet50_qiu/fullmodel_best.onnx --segment ../weights/segment_model.onnx"
            echo ""
            echo "  # 指定使用GPU 0和1"
            echo "  $0 --classifier ../weights/resnet50_qiu/fullmodel_best.onnx --gpu 0,1"
            exit 1
            ;;
    esac
done

# 检查至少配置了一个模型
if [ -z "$CLASSIFIER_MODEL_PATH" ] && [ -z "$SEGMENT_MODEL_PATH" ]; then
    echo "❌ 错误: 至少需要配置一个模型（分类或分割）"
    echo "使用 --classifier 指定分类模型，或使用 --segment 指定分割模型"
    exit 1
fi

# 检查模型文件是否存在
if [ -n "$CLASSIFIER_MODEL_PATH" ] && [ ! -f "$CLASSIFIER_MODEL_PATH" ]; then
    echo "❌ 错误: 分类模型文件不存在: $CLASSIFIER_MODEL_PATH"
    exit 1
fi

if [ -n "$SEGMENT_MODEL_PATH" ] && [ ! -f "$SEGMENT_MODEL_PATH" ]; then
    echo "❌ 错误: 分割模型文件不存在: $SEGMENT_MODEL_PATH"
    exit 1
fi

# 确定任务类型
if [ -n "$CLASSIFIER_MODEL_PATH" ] && [ -n "$SEGMENT_MODEL_PATH" ]; then
    # 同时配置了分类和分割模型，使用 cv
    TASK_TYPE="cv"
elif [ -n "$CLASSIFIER_MODEL_PATH" ]; then
    # 只配置了分类模型
    TASK_TYPE="classification"
elif [ -n "$SEGMENT_MODEL_PATH" ]; then
    # 只配置了分割模型
    TASK_TYPE="segmentation"
fi

CONTAINER_NAME="devdeploy-cv-${TASK_TYPE}-${PORT}"
COMPOSE_FILE="docker-compose.${TASK_TYPE}.${PORT}.yml"
# 固定镜像名称，所有容器共用同一个镜像
IMAGE_NAME="devdeploy-cv-api:latest"

# 模型路径转换为容器内路径
# 处理 ../weights/ 路径，转换为容器内的 /app/weights/ 路径
convert_to_container_path() {
    local model_path=$1
    if [[ "${model_path}" == ../weights/* ]]; then
        # 去掉 ../ 前缀，保留 weights/ 后面的部分
        RELATIVE_PATH="${model_path#../}"
        echo "/app/${RELATIVE_PATH}"
    elif [[ "${model_path}" == weights/* ]]; then
        # 如果已经是 weights/ 开头，直接拼接
        echo "/app/${model_path}"
    else
        # 其他情况，直接拼接（保持向后兼容）
        echo "/app/${model_path}"
    fi
}

CLASSIFIER_MODEL_PATH_IN_CONTAINER=""
if [ -n "$CLASSIFIER_MODEL_PATH" ]; then
    CLASSIFIER_MODEL_PATH_IN_CONTAINER=$(convert_to_container_path "$CLASSIFIER_MODEL_PATH")
fi

SEGMENT_MODEL_PATH_IN_CONTAINER=""
if [ -n "$SEGMENT_MODEL_PATH" ]; then
    SEGMENT_MODEL_PATH_IN_CONTAINER=$(convert_to_container_path "$SEGMENT_MODEL_PATH")
fi

echo "🚀 启动DevDeploy CV API服务（CCR版本）..."
echo ""
echo "配置信息:"
echo "  - 任务类型: ${TASK_TYPE}"
echo "  - 镜像名称: ${IMAGE_NAME}"
echo "  - 构建方式: 基于腾讯云CCR基础镜像"
if [ -n "$CLASSIFIER_MODEL_PATH" ]; then
    echo "  - 分类模型路径: ${CLASSIFIER_MODEL_PATH}"
    echo "  - 分类模型容器内路径: ${CLASSIFIER_MODEL_PATH_IN_CONTAINER}"
else
    echo "  - 分类模型: 未配置"
fi
if [ -n "$SEGMENT_MODEL_PATH" ]; then
    echo "  - 分割模型路径: ${SEGMENT_MODEL_PATH}"
    echo "  - 分割模型容器内路径: ${SEGMENT_MODEL_PATH_IN_CONTAINER}"
else
    echo "  - 分割模型: 未配置"
fi
echo "  - 端口: ${PORT}"
echo "  - 容器名称: ${CONTAINER_NAME}"
echo "  - Compose文件: ${COMPOSE_FILE}"
echo ""

# 检查Docker是否运行
if ! docker info > /dev/null 2>&1; then
    echo "❌ Docker未运行，请先启动Docker服务"
    exit 1
fi

echo "✅ 模型文件检查通过"

# 检查端口是否被占用
if netstat -tuln 2>/dev/null | grep -q ":${PORT}"; then
    echo "❌ 端口 ${PORT} 已被占用"
    echo "请使用其他端口"
    exit 1
fi

# 检查容器是否已存在
if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo "⚠️  容器 ${CONTAINER_NAME} 已存在"
    read -p "是否要停止并删除旧容器? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "正在停止旧容器..."
        docker-compose -f ${COMPOSE_FILE} down 2>/dev/null
        docker stop ${CONTAINER_NAME} > /dev/null 2>&1
        docker rm ${CONTAINER_NAME} > /dev/null 2>&1
    else
        echo "❌ 取消启动"
        exit 1
    fi
fi

# SSH端口计算
SSH_HOST_PORT=$((4022 + PORT - 8796))
SSH_CONTAINER_PORT=22

# 模板文件路径（使用标准模板，CCR版本和标准版本配置相同）
COMPOSE_TEMPLATE="docker-compose.cv.template.yml"

# 检查模板文件是否存在
if [ ! -f "${COMPOSE_TEMPLATE}" ]; then
    echo "❌ 模板文件不存在: ${COMPOSE_TEMPLATE}"
    echo "请确保模板文件在正确位置"
    exit 1
fi

# 检查Dockerfile.ccr是否存在
if [ ! -f "Dockerfile.ccr" ]; then
    echo "❌ Dockerfile.ccr 不存在"
    echo "请确保 Dockerfile.ccr 文件在正确位置"
    exit 1
fi

# 创建日志目录（如果不存在）
LOG_DIR="./logs"
if [ ! -d "$LOG_DIR" ]; then
    echo "📁 创建日志目录: ${LOG_DIR}"
    mkdir -p "$LOG_DIR"
fi

# 解析GPU配置
USE_GPU="false"
GPU_DEVICE_IDS="[]"
CUDA_VISIBLE_DEVICES=""

# 检查模板文件是否支持GPU
if grep -q "driver: nvidia" "${COMPOSE_TEMPLATE}" || grep -q "device_ids:" "${COMPOSE_TEMPLATE}"; then
    USE_GPU="true"
    
    # 函数：将GPU ID字符串转换为JSON数组格式和CUDA_VISIBLE_DEVICES
    # 输入: "0 1 2" 或 "0,1,2" (宿主机GPU编号)
    # 输出: GPU_DEVICE_IDS="[\"0\", \"1\", \"2\"]" (宿主机GPU编号，用于Docker device_ids)
    #      CUDA_VISIBLE_DEVICES="0,1,2" (容器内GPU编号，从0开始，因为Docker会重新映射)
    # 重要说明：
    # - Docker的device_ids指定宿主机GPU编号，容器内会重新映射为从0开始
    # - 例如：device_ids: ["1"] → 容器内只有cuda:0 → CUDA_VISIBLE_DEVICES=0
    # - 例如：device_ids: ["0","1"] → 容器内有cuda:0,cuda:1 → CUDA_VISIBLE_DEVICES=0,1
    parse_gpu_ids() {
        local input_ids="$1"
        # 清理输入：去除空格，统一为空格分隔
        local clean_ids=$(echo "$input_ids" | tr -d ' ' | tr ',' ' ' | xargs)
        
        if [ -z "$clean_ids" ]; then
            return 1
        fi
        
        # 验证并收集有效的GPU ID（宿主机编号）
        local host_gpu_ids=""
        local json_array="["
        local first=true
        local count=0
        
        for id in $clean_ids; do
            if [[ "$id" =~ ^[0-9]+$ ]]; then
                if [ "$first" = true ]; then
                    json_array="${json_array}\"${id}\""
                    host_gpu_ids="${id}"
                    first=false
                else
                    json_array="${json_array}, \"${id}\""
                    host_gpu_ids="${host_gpu_ids} ${id}"
                fi
                count=$((count + 1))
            fi
        done
        
        json_array="${json_array}]"
        
        if [ -n "$host_gpu_ids" ]; then
            # 设置全局变量（不使用local）
            GPU_DEVICE_IDS="$json_array"  # 宿主机GPU编号，用于Docker device_ids
            
            # CUDA_VISIBLE_DEVICES应该设置为容器内的GPU编号（从0开始）
            # 因为Docker的device_ids会将宿主机GPU重新映射为容器内从0开始编号
            # 例如：如果device_ids指定了2个GPU，容器内就是cuda:0和cuda:1
            local container_gpu_ids=""
            for i in $(seq 0 $((count - 1))); do
                if [ -z "$container_gpu_ids" ]; then
                    container_gpu_ids="$i"
                else
                    container_gpu_ids="${container_gpu_ids},$i"
                fi
            done
            
            CUDA_VISIBLE_DEVICES="$container_gpu_ids"
            return 0
        fi
        
        return 1
    }
    
    # 函数：从模板文件提取GPU ID
    extract_gpu_from_template() {
        if grep -q "device_ids: {{GPU_DEVICE_IDS}}" "${COMPOSE_TEMPLATE}"; then
            # 使用占位符，返回默认值
            echo "0"
            return 0
        fi
        
        # 尝试提取实际的GPU设备ID
        local gpu_line=$(grep "device_ids:" "${COMPOSE_TEMPLATE}" | head -1)
        if [[ "$gpu_line" =~ device_ids:[[:space:]]*\[(.*)\] ]]; then
            # 清理引号和空格，返回空格分隔的ID
            echo "${BASH_REMATCH[1]}" | tr -d "'\" " | tr ',' ' ' | xargs
            return 0
        fi
        
        return 1
    }
    
    # 函数：验证GPU设备是否存在
    validate_gpu_devices() {
        if [ -z "$CUDA_VISIBLE_DEVICES" ]; then
            return 0
        fi
        
        if ! command -v nvidia-smi &> /dev/null; then
            echo "⚠️  警告: nvidia-smi不可用，跳过GPU设备验证"
            return 0
        fi
        
        local gpu_count=$(nvidia-smi --list-gpus 2>/dev/null | wc -l)
        if [ -z "$gpu_count" ] || [ "$gpu_count" -eq 0 ]; then
            echo "❌ 错误: 系统中未检测到GPU设备"
            exit 1
        fi
        
        local invalid_gpus=""
        for gpu_id in $(echo "$CUDA_VISIBLE_DEVICES" | tr ',' ' '); do
            if [ "$gpu_id" -ge "$gpu_count" ]; then
                invalid_gpus="${invalid_gpus:+$invalid_gpus, }$gpu_id"
            fi
        done
        
        if [ -n "$invalid_gpus" ]; then
            echo "❌ 错误: 指定的GPU设备不存在: $invalid_gpus"
            echo "   系统中可用的GPU设备: 0 到 $((gpu_count - 1))"
            echo "   请使用 --gpu 参数指定正确的GPU设备ID，例如: --gpu 0"
            exit 1
        fi
    }
    
    # 确定使用的GPU ID（优先级：命令行参数 > 模板文件 > 默认值0）
    if [ -n "$GPU_IDS_ARG" ]; then
        # 优先使用命令行指定的GPU ID
        if ! parse_gpu_ids "$GPU_IDS_ARG"; then
            echo "⚠️  警告: --gpu 参数无效，使用默认GPU配置"
            parse_gpu_ids "0"
        fi
    else
        # 从模板文件提取，如果失败则使用默认值
        local template_gpu_ids=$(extract_gpu_from_template)
        parse_gpu_ids "${template_gpu_ids:-0}" || parse_gpu_ids "0"
    fi
    
    # 验证GPU设备
    validate_gpu_devices
fi

echo "📝 生成 docker-compose 配置文件: ${COMPOSE_FILE}..."
echo "  - GPU配置: USE_GPU=${USE_GPU}, GPU_DEVICE_IDS=${GPU_DEVICE_IDS}, CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}"
# 使用 sed 替换模板中的占位符
# 如果模型路径为空，使用空字符串
sed -e "s|{{IMAGE_NAME}}|${IMAGE_NAME}|g" \
    -e "s|{{CONTAINER_NAME}}|${CONTAINER_NAME}|g" \
    -e "s|{{PORT}}|${PORT}|g" \
    -e "s|{{SSH_HOST_PORT}}|${SSH_HOST_PORT}|g" \
    -e "s|{{MODEL_PATH_IN_CONTAINER}}|${CLASSIFIER_MODEL_PATH_IN_CONTAINER}|g" \
    -e "s|{{SEGMENT_MODEL_PATH_IN_CONTAINER}}|${SEGMENT_MODEL_PATH_IN_CONTAINER}|g" \
    -e "s|{{USE_GPU}}|${USE_GPU}|g" \
    -e "s|{{GPU_DEVICE_IDS}}|${GPU_DEVICE_IDS}|g" \
    -e "s|{{CUDA_VISIBLE_DEVICES}}|${CUDA_VISIBLE_DEVICES}|g" \
    "${COMPOSE_TEMPLATE}" > "${COMPOSE_FILE}"

echo "✅ 配置文件已生成"

# 检查Docker镜像是否存在，如果不存在则构建
if ! docker images --format '{{.Repository}}:{{.Tag}}' | grep -q "^${IMAGE_NAME}$"; then
    echo "🔨 构建Docker镜像: ${IMAGE_NAME}（使用Dockerfile.ccr）..."
    echo "📦 基础镜像: ccr.ccs.tencentyun.com/ai_plateform/devdeploy-cv-api:latest"
    docker build -t ${IMAGE_NAME} -f Dockerfile.ccr .
    if [ $? -ne 0 ]; then
        echo "❌ 镜像构建失败"
        rm -f ${COMPOSE_FILE}
        exit 1
    fi
    echo "✅ 镜像构建完成"
else
    echo "✅ 镜像已存在: ${IMAGE_NAME}，跳过构建"
fi

# 启动服务
echo "🚀 启动服务..."
docker-compose -f ${COMPOSE_FILE} up -d

# 等待服务启动
echo "⏳ 等待服务启动..."
sleep 10

# 检查服务状态
if docker ps | grep -q "${CONTAINER_NAME}"; then
    echo "✅ 服务启动成功！"
    echo ""
    echo "📋 服务信息："
    echo "  - API地址: http://localhost:${PORT}"
    echo "  - 健康检查: http://localhost:${PORT}/health"
    echo "  - API文档: http://localhost:${PORT}/docs"
    echo "  - 配置信息: http://localhost:${PORT}/config"
    echo "  - SSH端口: ${SSH_HOST_PORT}"
    echo ""
    echo "📝 日志文件位置："
    echo "  - 日志文件: ./logs/cv_api_port_${PORT}.log"
    echo "  - 实时查看: tail -f ./logs/cv_api_port_${PORT}.log"
    echo ""
    echo "🔍 查看容器日志: docker-compose -f ${COMPOSE_FILE} logs -f"
    echo "🛑 停止服务: docker-compose -f ${COMPOSE_FILE} down"
    echo "🔄 重启服务: docker-compose -f ${COMPOSE_FILE} restart"
else
    echo "❌ 服务启动失败"
    echo "查看日志: docker-compose -f ${COMPOSE_FILE} logs"
    exit 1
fi
